import pandas as pd
import os
import numpy as np


class Data_consentrate():
    
    #将数据格式整理为dataframe，列名为股票名
    def __init__(self,  root=r'data/stock_data/daily_price', thred=1000):
        self.thred = thred
        self.root = root    
        code_list = os.listdir(self.root) 
        code_list.sort()
        self.code_list = code_list
        # print(f'股票顺序{self.code_list}')
        _root = '/'.join(root.split('/')[:-1])
        self.target_root = os.path.join(_root, f"consentrate_{root.split('/')[-1]}")
        if not os.path.exists(self.target_root):
            os.makedirs(self.target_root) 
        
    def consentrate(self, key): #['close',volume]
        print(f'开始集中数据:{key}')
        drop_list = []
        data_lst = []
        column_lst = []
        for i in self.code_list:
            root_i = os.path.join(self.root,i)
            data_i = pd.read_pickle(root_i)
            length_i = len(data_i)
            if length_i <=self.thred:
                print(f'去除新股：{i}, 数据长度{length_i}')
                drop_list.append(i)
                continue
            name_i = i[:-9]
            if key == 'lagret':
                data_i = self.lagret(data_i)
            elif key == 'vol':
                data_i = self.std10(data_i)
            else:     
                data_i = data_i[[key]]
            data_lst.append(data_i)
            column_lst.append(name_i)
        data = pd.concat(data_lst, axis=1)
        data.columns = column_lst
        data.to_pickle(f'{self.target_root}/{key}.pkl.gzip')
        drop_code = pd.DataFrame(drop_list,columns=['drop_code']) 
        drop_code.to_csv(f'{self.target_root}/drop_code.csv')
        
    def lagret(self, data_i):
        return data_i['close'] / data_i['close'].shift(1)-1

    def std10(self, data_i):
        lagret = np.log(data_i['close'] / data_i['close'].shift(1))
        return lagret.rolling(100, min_periods=1).std()

class consentrate_price_value():
    def __init__(self, price_root=None, value_root=None, thred=1000) -> None:
        self.thred = thred
        self.price_root = r'data\stock_data\daily\price' if price_root is None else price_root
        self.value_root = r'data\stock_data\daily\market_value' if price_root is None else price_root
        self.save_root = os.path.join(os.path.split(self.price_root)[0],'consentrate_price')
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.code_list = sorted(os.listdir(self.price_root))

    def consentrate(self, key): 
        print(f'开始集中数据:{key}')
        drop_list = []
        data_lst = []
        column_lst = []
        for i in self.code_list:
            price_root_i = os.path.join(self.price_root,i)
            value_root_i = os.path.join(self.value_root,i)
            price_i = pd.read_pickle(price_root_i)
            value_i = pd.read_pickle(value_root_i)
            data_i = pd.concat([price_i, value_i], axis=1)
            length_i = len(data_i)
            if length_i <=self.thred:
                print(f'去除新股：{i}, 数据长度{length_i}')
                drop_list.append(i)
                continue
            name_i = i[:-9]
            if key == 'lagret':
                data_i = self.lagret(data_i)
            elif key == 'vol':
                data_i = self.std10(data_i)
            else:     
                data_i = data_i[[key]]
            data_lst.append(data_i)
            column_lst.append(name_i)
        data = pd.concat(data_lst, axis=1)
        data.columns = column_lst
        data.to_pickle(f'{self.save_root}/{key}.pkl.gzip')
        drop_code = pd.DataFrame(drop_list,columns=['drop_code']) 
        drop_code.to_csv(f'{self.save_root}/drop_code.csv')

        
    def lagret(self, data_i):
        return np.log(data_i['close'] / data_i['close'].shift(1))

    def std10(self, data_i):
        lagret = np.log(data_i['close'] / data_i['close'].shift(1))
        return lagret.rolling(100, min_periods=1).std()

    def consentrate_columns(self):
        columns = ['open', 'close', 'high', 'low', 'volume', 'amount', 'pe', 'pe_ttm',
       'pb', 'ps', 'ps_ttm', 'dv_ratio', 'dv_ttm', 'total_mv', 'lagret', 'vol']
        for i in columns:
            self.consentrate(i)


# cc = consentrate_price_value()
# cc.consentrate_columns()